# FLEXPART interface module for reading FLEXPART output headers
#
# This module provides functions to read FLEXPART header files and extract
# grid information needed for initialization and domain validation.

module FlexpartIO

# Define constants locally to avoid circular dependencies during testing
const R_earth = 6.371e6  # Earth radius in meters
const π = 3.14159265358979323846

"""
    FlexpartHeader

Struct containing FLEXPART grid information read from header files.
Fields correspond to Fortran variables in mod_var.f90:
- nxgrid, nygrid: grid dimensions
- llx, lly: lower left corner coordinates
- dx, dy: grid spacing
- nxshift: longitudinal shift for global grids
- outheight: height levels (meters)
- nzgrid: number of vertical levels
"""
Base.@kwdef struct FlexpartHeader
    nxgrid::Int = 0
    nygrid::Int = 0
    llx::Float64 = 0.0
    lly::Float64 = 0.0
    dx::Float64 = 1.0
    dy::Float64 = 1.0
    nxshift::Int = 0
    outheight::Vector{Float64} = Float64[]
    nzgrid::Int = 0
    trajdays::Int = 0
end

"""
    read_header(path::AbstractString) -> FlexpartHeader

Read FLEXPART header file and extract grid information.

Based on the Fortran subroutine read_header in mod_flexpart.f90.
Reads the unformatted binary header file to get grid dimensions and parameters.

# Arguments
- `path`: Path to the FLEXPART header file

# Returns
- `FlexpartHeader`: Struct containing grid information

# Notes
- For global domains (longitude span = 360°), calculates xshift relative to 180°W
- This function only reads the grid information needed for initialization
- The full header contains additional information (species, releases, etc.) that
  is not needed for the current implementation
"""
function read_header(path::AbstractString)
    if !isfile(path)
        error("Cannot find FLEXPART header file: $path")
    end

    @warn "Reading binary FLEXPART header - this is experimental and may not work correctly. Prefer header_txt files when available."

    # Read header from unformatted binary file
    # Based on structure in mod_flexpart.f90:read_header
    open(path, "r") do io
        # Skip initial data we don't need for grid info
        # Read: ibdate, ibtime, flexversion
        seek(io, 4)  # Skip record length marker
        ibdate = read(io, Int32)
        ibtime = read(io, Int32)
        flexversion = strip(String(read(io, 29)))
        seek(io, position(io) + 4)  # Skip trailing record marker

        # Read: loutstep, loutaver, loutsample
        seek(io, position(io) + 4)  # Skip record length marker
        loutstep = read(io, Int32)
        loutaver = read(io, Int32)
        loutsample = read(io, Int32)
        seek(io, position(io) + 4)  # Skip trailing record marker

        # Read grid setup: bndx, bndy, numx, numy, delx, dely
        seek(io, position(io) + 4)  # Skip record length marker
        bndx = read(io, Float32)
        bndy = read(io, Float32)
        numx = read(io, Int32)
        numy = read(io, Int32)
        delx = read(io, Float32)
        dely = read(io, Float32)
        seek(io, position(io) + 4)  # Skip trailing record marker

        # Read vertical grid: nzgrid, outheight(1:nzgrid)
        seek(io, position(io) + 4)  # Skip record length marker
        nzgrid = read(io, Int32)
        outheight = Vector{Float64}(undef, nzgrid)
        for i = 1:nzgrid
            outheight[i] = Float64(read(io, Float32))
        end
        seek(io, position(io) + 4)  # Skip trailing record marker

        # Calculate xshift for global domains (matching Fortran logic)
        xshift = 0
        if abs(Float64(numx) * Float64(delx) - 360.0) < 1e-6
            # Global domain - check offset relative to 180W
            xshift = Int(floor((Float64(bndx) + 180.0) / Float64(delx)))
            bndx = -180.0
        end

        return FlexpartHeader(
            nxgrid = Int(numx),
            nygrid = Int(numy),
            llx = Float64(bndx),
            lly = Float64(bndy),
            dx = Float64(delx),
            dy = Float64(dely),
            nxshift = xshift,
            outheight = outheight,
            nzgrid = Int(nzgrid)
        )
    end
end

"""
    read_header_text(path::AbstractString) -> FlexpartHeader

Read FLEXPART header from text format (header_txt files).

This is an alternative reader for the text format header files that
are sometimes available alongside the binary headers.

# Arguments
- `path`: Path to the FLEXPART header_txt file

# Returns
- `FlexpartHeader`: Struct containing grid information
"""
function read_header_text(path::AbstractString)
    if !isfile(path)
        error("Cannot find FLEXPART header text file: $path")
    end

    lines = readlines(path)
    trajdays = 0

    # Find grid setup line (should be after the comment about grid setup)
    grid_line = ""
    for (i, line) in enumerate(lines)
        # Look for comment about grid setup, then take the next data line
        if occursin("outlon0", line) && i < length(lines)
            # Take the next non-comment line
            for j = i+1:length(lines)
                if !startswith(strip(lines[j]), "#") && !isempty(strip(lines[j]))
                    grid_line = lines[j]
                    break
                end
            end
            break
        end
        # Alternative: look for lines with 6 numeric values
        if !startswith(strip(line), "#") && !isempty(strip(line))
            parts = split(strip(line))
            if length(parts) >= 6 && all(x -> tryparse(Float64, x) !== nothing, parts[1:6])
                grid_line = line
                break
            end
        end
    end

    if isempty(grid_line)
        error("Could not find grid setup line in header file")
    end

    # Parse grid parameters: outlon0, outlat0, numxgrid, numygrid, dxout, dyout
    parts = split(strip(grid_line))
    if length(parts) < 6
        error("Invalid grid setup line format: found $(length(parts)) parts, need 6")
    end

    bndx = parse(Float64, parts[1])
    bndy = parse(Float64, parts[2])
    numx = parse(Int, parts[3])
    numy = parse(Int, parts[4])
    delx = parse(Float64, parts[5])
    dely = parse(Float64, parts[6])

    # Parse vertical levels - look for comment about numzgrid, then data line
    nzgrid = 0
    outheight = Float64[]

    for (i, line) in enumerate(lines)
        if occursin("numzgrid", line) && i < length(lines)
            # Take the next non-comment line
            for j = i+1:length(lines)
                if !startswith(strip(lines[j]), "#") && !isempty(strip(lines[j]))
                    parts = split(strip(lines[j]))
                    if length(parts) >= 1
                        nzgrid = parse(Int, parts[1])
                        # Heights should be the remaining values on the same line
                        if length(parts) > 1
                            outheight = [parse(Float64, parts[k]) for k = 2:min(length(parts), nzgrid+1)]
                        end
                    end
                    break
                end
            end
            break
        end
    end

    # Calculate xshift for global domains
    xshift = 0
    if abs(Float64(numx) * delx - 360.0) < 1e-6
        # Global domain - check offset relative to 180W
        xshift = Int(floor((bndx + 180.0) / delx))
        bndx = -180.0
    end

    # Extract backtrajectory length (age class) if available
    for (i, line) in enumerate(lines)
        if occursin("information on age class", lowercase(line)) && i < length(lines)
            for j = i+1:length(lines)
                s = strip(lines[j])
                isempty(s) && continue
                startswith(s, "#") && continue
                parts = split(s)
                if length(parts) >= 2
                    age_seconds = try
                        parse(Float64, parts[2])
                    catch
                        0.0
                    end
                    if age_seconds > 0
                        trajdays = Int(round(age_seconds / 86400.0))
                    end
                end
                break
            end
            break
        end
    end

    return FlexpartHeader(
        nxgrid = numx,
        nygrid = numy,
        llx = bndx,
        lly = bndy,
        dx = delx,
        dy = dely,
        nxshift = xshift,
        outheight = outheight,
        nzgrid = nzgrid,
        trajdays = trajdays
    )
end

end # module
